// Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "dali/operators/random/noise/shot_noise.h"
#include "dali/operators/random/rng_base_cpu.h"
#include "dali/pipeline/operator/arg_helper.h"

namespace dali {

DALI_SCHEMA(noise__Shot)
    .DocStr(R"code(Applies shot noise to the input.

The shot noise is generated by applying the following formula::

    output[:] = poisson_dist(max(0, input[:] / factor)) * factor) if factor != 0
    output[:] = input[:]                                          if factor == 0

where ``poisson_dist`` represents a poisson distribution.

Shot noise is a noise that's present in data generated by a Poisson process, like
registering photons by an image sensor. This operator simulates the data
acquisition process where each event increases the output value by
`factor` and the input tensor contains the expected values of corresponding
output points. For example, a `factor` of 0.1 means that 10 events are
needed to increase the output value by 1, while a factor of 10 means that
a single event increases the output by 10. The output values are quantized
to multiples of `factor`. The larger the factor, the more noise is present in
the output. A factor of 0 makes this an identity operation.

The shape and data type of the output will match the input.
)code")
    .NumInput(1)
    .NumOutput(1)
    .AddRandomSeedArg()
    .AddRandomStateArg()
    .AddOptionalArg<float>("factor",
      R"code(Factor parameter.)code",
     20.0f, true);

DALI_REGISTER_OPERATOR(noise__Shot, ShotNoise<CPUBackend>, CPU);

}  // namespace dali
